package edu.nepu.flink.api.watermark;

import edu.nepu.flink.api.bean.WaterSensor;
import org.apache.flink.api.common.eventtime.*;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;

import java.time.Duration;

/**
 * @Date 2024/3/1 14:54
 * @Created by chenshuaijun
 */
public class SelfDefineWatermarkGeneStrategy {

    public static void main(String[] args) {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        env.setParallelism(1);

        SingleOutputStreamOperator<WaterSensor> streamSource = env.socketTextStream("hadoop102", 9999).map(new MapFunction<String, WaterSensor>() {
            @Override
            public WaterSensor map(String value) throws Exception {
                String[] split = value.split(" ");
                return new WaterSensor(split[0],Long.parseLong(split[1]),Integer.parseInt(split[2]));
            }
        });
        SingleOutputStreamOperator<WaterSensor> watermarkStream = streamSource.assignTimestampsAndWatermarks(new MyWaterStage());

        watermarkStream.keyBy(WaterSensor::getId)
                       .window(TumblingEventTimeWindows.of(Time.seconds(10)))
                       .reduce(new ReduceFunction<WaterSensor>() {
                           @Override
                           public WaterSensor reduce(WaterSensor value1, WaterSensor value2) throws Exception {
                               return new WaterSensor(value1.id, value2.ts, value1.vc + value2.vc);
                           }
                       }).print();
    }

    static class MyWaterStage implements WatermarkStrategy<WaterSensor>{

        @Override
        public WatermarkGenerator<WaterSensor> createWatermarkGenerator(WatermarkGeneratorSupplier.Context context) {
            return new MyGeneWaterMarkPeriodic(2000);
        }

        @Override
        public TimestampAssigner<WaterSensor> createTimestampAssigner(TimestampAssignerSupplier.Context context) {
            return new SerializableTimestampAssigner<WaterSensor>() {
                @Override
                public long extractTimestamp(WaterSensor element, long recordTimestamp) {
                    return element.getTs() * 1000;
                }
            };
        }
    }

    /**
     * 这里定义的是周期性触发的
     */
    static class MyGeneWaterMarkPeriodic implements WatermarkGenerator<WaterSensor>{
        // 定义maxTs：指代的是当前最大的时间戳
        private long maxTs = 0L;
        // 指定延迟时间
        private long lateTs = 0L;

        public MyGeneWaterMarkPeriodic(long lateTs) {
            this.lateTs = lateTs;
            this.maxTs = Long.MIN_VALUE + lateTs + 1;
        }

        @Override
        public void onEvent(WaterSensor event, long eventTimestamp, WatermarkOutput output) {
            maxTs = Math.max(maxTs,eventTimestamp);
        }

        @Override
        public void onPeriodicEmit(WatermarkOutput output) {
            output.emitWatermark(new Watermark(maxTs - lateTs - 1));
        }
    }

    /**
     * 这里定义的是间歇性触发的，也就是每条数据到来的时候就触发一次
     */
    static class MyGeneWaterMark implements WatermarkGenerator<WaterSensor>{
        // 定义maxTs：指代的是当前最大的时间戳
        private long maxTs = 0L;
        // 指定延迟时间
        private long lateTs = 0L;

        public MyGeneWaterMark(long lateTs) {
            this.lateTs = lateTs;
            this.maxTs = Long.MIN_VALUE + lateTs + 1;
        }

        @Override
        public void onEvent(WaterSensor event, long eventTimestamp, WatermarkOutput output) {
            maxTs = Math.max(maxTs,eventTimestamp);
            // 这个是发射watermark，如果是周期性的触发就将其写到下面的onPeriodicEmit方法中，如果是间歇性的，就写到这个方法中，这样来一条数据就会触发一次
            output.emitWatermark(new Watermark(maxTs - lateTs - 1));
        }

        @Override
        public void onPeriodicEmit(WatermarkOutput output) {

        }
    }
}
